package cn.kduck.code.generate.service;


import cn.kduck.code.generate.json.AiChatRequest.ChatMessage;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

public class MessageChatMemory {

    Map<String, List<ChatMessage>> conversationHistory = new ConcurrentHashMap<>();

    private String topPrompt;

    public void add(String conversationId, List<ChatMessage> messages) {
        this.conversationHistory.putIfAbsent(conversationId, new ArrayList<>());
        this.conversationHistory.get(conversationId).addAll(messages);
    }

    public List<ChatMessage> get(String conversationId, int lastN) {
        List<ChatMessage> all = this.conversationHistory.get(conversationId);
        List<ChatMessage> lastMessage = all != null ? all.stream().skip(Math.max(0, all.size() - lastN)).collect(Collectors.toList()) : new ArrayList<>();
        if(!lastMessage.isEmpty()){
            lastMessage.set(0,new ChatMessage("user",this.topPrompt));
        }
        return lastMessage;
    }

    public void setTopPrompt(String topPrompt) {
        this.topPrompt = topPrompt;
    }

    public void clear(String conversationId) {
        this.conversationHistory.remove(conversationId);
    }
}
